import * as tf from '@tensorflow/tfjs'
import * as tfvis from '@tensorflow/tfjs-vis'
import { MnistData } from './data.js'


window.onload = async () => {
	const data = new MnistData()
	await data.load()

	// data 的形状是 784 = 28 * 28 * 1 因为每张图片像素
	const examples = data.nextTestBatch(20)
	console.log(examples)
	const surface = tfvis.visor().surface({name: '输入示例', tab: 'Tab1'})
	for (let i = 0; i < examples.labels.shape[0]; i ++) {
		/*
		* 执行提供的函数fn，在它被执行后，清除除fn返回的外fn分配的所有中间张量。
		* fn不能返回Promise(不允许使用异步函数)。返回的结果可以是一个复杂对象。
		* 使用此方法有助于避免内存泄漏。通常，在tf.tidy()中包装对操作的调用以实现自动内存清理。
		* */
		const imageTensor = tf.tidy(() => {
			return examples.xs
						.slice([i, 0], [1, 784])
						.reshape([28, 28, 1])
		})


		const canvas = document.createElement('canvas')
		canvas.height = 28
		canvas.width = 28
		await tf.browser.toPixels(imageTensor, canvas)
		surface.drawArea.append(canvas)
	}

	// 定义模型
	let model = null
	// try {
	// 	model = await tf.loadLayersModel('localstorage://mnist-model')
	// } catch (e) {}
	if (!model) {
		model = tf.sequential()

		// 添加卷积层
		/*
		* inputShape: 输入形状
		* kernelSize: 卷积核 大小
		* filters: 多个组过滤器
		* strides: 步长
		* activation: relu 激活函数，过滤 负数 f(x) = max(0, x)
		* kernelInitializer: 卷积核初始化方法
		* */
		model.add(tf.layers.conv2d({
			inputShape: [28, 28, 1],
			kernelSize: 5,
			filters: 8,
			strides: 1,
			activation: 'relu',
			kernelInitializer: 'varianceScaling'
		}))

		// 添加池化层, 最大池化层
		model.add(tf.layers.maxPool2d({
			poolSize: [2, 2],
			strides: [2, 2]
		}))

		// 重复卷积和池化
		model.add(tf.layers.conv2d({
			kernelSize: 5,
			filters: 16,
			strides: 1,
			activation: 'relu',
			kernelInitializer: 'varianceScaling'
		}))

		model.add(tf.layers.maxPool2d({
			poolSize: [2, 2],
			strides: [2, 2]
		}))

		// 拉平
		model.add(tf.layers.flatten())

		/* 全连接层
		 * units 神经元数
		 * softmax : 使10个输出 总和为1
		 */
		model.add(tf.layers.dense({
			units: 10,
			activation: 'softmax',
			kernelInitializer: 'varianceScaling'
		}))

		/*
		 * 设置损失函数: 交叉熵 categoricalCrossentropy 是对数函数的多分类版本
		 * 优化器adam(学习率)
		 */
		model.compile({
			loss: "categoricalCrossentropy",
			optimizer: tf.train.adam(),
			metrics: ['accuracy']
		})

		const [trainX, trainY] = tf.tidy(() => {
			const d = data.nextTrainBatch(2000)
			return [
				d.xs.reshape([2000, 28, 28, 1]),
				d.labels
			]
		})

		const [testX, testY] = tf.tidy(() => {
			const d = data.nextTrainBatch(200)
			return [
				d.xs.reshape([200, 28, 28, 1]),
				d.labels
			]
		})

		/** fit 开始学习任务
		 * batchSize 取样数
		 * epochs 对训练数据数组进行迭代的次数。
		 * loss: 训练集损失
		 * val_loss: 验证集损失
		 * acc: 训练集准确度
		 * val_acc: 验证集准确度
		 * callbacks: 只展示 onEpochEnd 的曲线图
		 */
		model.fit(trainX, trainY, {
			batchSize: 500,
			epochs: 50,
			validationData: [testX, testY],
			callbacks: tfvis.show.fitCallbacks({
				name: "训练过程",
				tab: 'Tab2'
			}, [
				'loss', 'val_loss', 'acc', 'val_acc'
			], {
				callbacks: ['onEpochEnd'] //onEpochEnd 结束时
			})
		})

		await model.save('localstorage://mnist-model');
	}

	const canvas = document.querySelector('canvas');
	canvas.addEventListener('mousemove', (e) => {
		if (e.buttons === 1) {
			const ctx = canvas.getContext('2d');
			ctx.fillStyle = 'rgb(255,255,255)';
			ctx.fillRect(e.offsetX, e.offsetY, 25, 25);
		}
	})

	window.clear = () => {
		const ctx = canvas.getContext('2d');
		ctx.fillStyle = 'rgb(0,0,0)';
		ctx.fillRect(0, 0, 300, 300);
	};

	clear();

	window.predict = () => {
		const input = tf.tidy(() => {

			/*
			* tf.browser.fromPixels 像素转 tensor
			* tf.image.resizeBilinear 改变形状
			* slice 降维置灰
			* toFloat， div 运用运算归一化
			* */
			const d = tf.image
			.resizeBilinear(tf.browser.fromPixels(canvas), [28, 28], true)
			.slice([0, 0, 0], [28, 28, 1]).toFloat().div(255).reshape([1, 28, 28, 1])
			d.print()
			console.log(d)



			return tf.image
			.resizeBilinear(tf.browser.fromPixels(canvas), [28, 28], true)
			.slice([0, 0, 0], [28, 28, 1])
			.toFloat()
			.div(255)
			.reshape([1, 28, 28, 1])
		})

		const output = model.predict(input);
		// argMax(n) 表示第 n 维最大值
		alert(`预测为 ${output.argMax(1).dataSync()[0]}`);
	}
}
